/*
 * Tencent is pleased to support the open source community by making Angel available.
 *
 * Copyright (C) 2017 THL A29 Limited, a Tencent company. All rights reserved.
 *
 * Licensed under the BSD 3-Clause License (the "License"); you may not use this file except in
 * compliance with the License. You may obtain a copy of the License at
 *
 * https://opensource.org/licenses/BSD-3-Clause
 *
 * Unless required by applicable law or agreed to in writing, software distributed under the License is
 * distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND,
 * either express or implied. See the License for the specific language governing permissions and
 * limitations under the License.
 */

package com.tencent.angel.master.task;

import com.tencent.angel.worker.task.TaskId;
import java.io.DataInputStream;
import java.io.DataOutputStream;
import java.io.IOException;
import java.util.Collection;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Map.Entry;
import java.util.Set;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReadWriteLock;
import java.util.concurrent.locks.ReentrantReadWriteLock;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;

/**
 * Task manager in master, it manages all tasks {@link com.tencent.angel.master.task.AMTask} in the
 * application.
 */
// TODO: 17/6/26 by zmyer
public class AMTaskManager {
    private static final Log LOG = LogFactory.getLog(AMTaskManager.class);
    /** task id to task map */
    private final HashMap<TaskId, AMTask> idToTaskMap;
    private final Lock readLock;
    private final Lock writeLock;

    public AMTaskManager() {
        idToTaskMap = new HashMap<>();
        ReadWriteLock readWriteLock = new ReentrantReadWriteLock();
        readLock = readWriteLock.readLock();
        writeLock = readWriteLock.writeLock();
    }

    /**
     * add the task to task map
     *
     * @param taskId the task id
     * @param task the task
     */
    public void putTask(TaskId taskId, AMTask task) {
        try {
            writeLock.lock();
            idToTaskMap.put(taskId, task);
        } finally {
            writeLock.unlock();
        }
    }

    /**
     * get the task corresponding to specified task id
     *
     * @param taskId task id
     * @return AMTask the task corresponding to specified task id
     */
    public AMTask getTask(TaskId taskId) {
        try {
            readLock.lock();
            return idToTaskMap.get(taskId);
        } finally {
            readLock.unlock();
        }
    }

    /**
     * get all tasks
     *
     * @return Collection<AMTask> all tasks
     */
    public Collection<AMTask> getTasks() {
        try {
            readLock.lock();
            Set<AMTask> cloneTasks = new HashSet<>();
            cloneTasks.addAll(idToTaskMap.values());
            return cloneTasks;
        } finally {
            readLock.unlock();
        }
    }

    /**
     * write all tasks state to a output stream
     *
     * @param output the output stream
     * @throws IOException io exception that throw
     */
    public void serialize(DataOutputStream output) throws IOException {
        try {
            readLock.lock();
            output.writeInt(idToTaskMap.size());
            for (Entry<TaskId, AMTask> entry : idToTaskMap.entrySet()) {
                output.writeInt(entry.getKey().getIndex());
                entry.getValue().serialize(output);
            }
        } finally {
            readLock.unlock();
        }
    }

    /**
     * read all tasks state from a input stream
     *
     * @param input the input stream
     * @throws IOException io exception that throw
     */
    public void deserialize(DataInputStream input) throws IOException {
        try {
            writeLock.lock();
            int size = input.readInt();
            for (int i = 0; i < size; i++) {
                TaskId taskId = new TaskId(input.readInt());
                AMTask task = new AMTask(taskId, null);
                task.deserialize(input);
                LOG.info("task deserialize=" + task);
                idToTaskMap.put(taskId, task);
            }
        } finally {
            writeLock.unlock();
        }
    }
}
